package raft

import (
	"slices"
	"time"
)

type InstallSnapshotArgs struct {
	Term              int
	LeaderId          int
	LastIncludedIndex int
	LastIncludedTerm  int
	Data              []byte
	// 实验要求在单个RPC中发送整个snapshot
	// Offset            int
	// Done              bool
}

type InstallSnapshotReply struct {
	Term int
}

// 这个函数没加锁
func (rf *Raft) abs2Relative(absIndex int) int {
	return absIndex - rf.lastIncludedIndex
}

// 这个函数没加锁
func (rf *Raft) relative2Abs(relativeIndex int) int {
	return relativeIndex + rf.lastIncludedIndex
}

func (rf *Raft) sendInstallSnapshot(server int, args *InstallSnapshotArgs, reply *InstallSnapshotReply) bool {
	ok := rf.peers[server].Call("Raft.InstallSnapshot", args, reply)
	return ok
}

// 没加锁，返回snapshot是否有更新
func (rf *Raft) updateSnapshot(lastIncludedIndex, lastIncludedTerm int, data []byte) bool {
	if lastIncludedIndex <= rf.lastIncludedIndex {
		return false
	}
	if relativeLastIncludedIndex := rf.abs2Relative(lastIncludedIndex); relativeLastIncludedIndex < len(rf.Log) {
		rf.Log = slices.Clone(rf.Log[relativeLastIncludedIndex:])
	} else {
		rf.Log = []Entry{{Term: lastIncludedTerm, Command: nil}}
	}
	rf.lastIncludedIndex = lastIncludedIndex
	rf.snapshot = data
	rf.persist()
	return true
}

func (rf *Raft) InstallSnapshot(args *InstallSnapshotArgs, reply *InstallSnapshotReply) {
	needPersist := false
	rf.mu.Lock()
	defer rf.mu.Unlock()
	if args.Term < rf.CurrentTerm {
		reply.Term = rf.CurrentTerm
		return
	}
	if args.Term > rf.CurrentTerm {
		rf.CurrentTerm = args.Term
		rf.becomeFollower()
		needPersist = true
	}
	if rf.updateSnapshot(args.LastIncludedIndex, args.LastIncludedTerm, args.Data) {
		go rf.startApplyService(0)
		// Debug(dLog, "S%d T%d, install snapshot index %d, term %d, len %d", rf.me, rf.CurrentTerm, rf.lastIncludedIndex, rf.Log[0].Term, len(args.Data))
	} else if needPersist {
		rf.persist()
	}
	reply.Term = rf.CurrentTerm
}

// 这里有锁；返回是否成功(能否下班)
func (rf *Raft) installSnapshotOnceTo(server int) bool {
	rf.mu.Lock()
	if rf.status != LEADER || rf.nextIndex[server] == rf.relative2Abs(len(rf.Log)) {
		rf.mu.Unlock()
		return true
	}
	if rf.nextIndex[server] > rf.lastIncludedIndex {
		rf.mu.Unlock()
		go rf.appendEntriesTo(server)
		return true
	}
	args := &InstallSnapshotArgs{
		Term:              rf.CurrentTerm,
		LeaderId:          rf.me,
		LastIncludedIndex: rf.lastIncludedIndex,
		LastIncludedTerm:  rf.Log[0].Term,
		Data:              make([]byte, len(rf.snapshot)),
	}
	copy(args.Data, rf.snapshot)
	// Debug(dLeader, "S%d T%d installSnapshotOnceTo to S%d, index %d term %d len %d", rf.me, rf.CurrentTerm, server, args.LastIncludedIndex, args.LastIncludedTerm, len(args.Data))
	rf.mu.Unlock()

	reply := &InstallSnapshotReply{}
	ok := rf.sendInstallSnapshot(server, args, reply)
	if !ok {
		return false
	}
	rf.mu.Lock()
	defer rf.mu.Unlock()
	if reply.Term > rf.CurrentTerm {
		rf.CurrentTerm = reply.Term
		rf.becomeFollower()
		rf.persist()
		return true
	}
	if rf.status != LEADER || args.Term < rf.CurrentTerm {
		return true
	}
	rf.matchIndex[server] = max(args.LastIncludedIndex, rf.matchIndex[server])
	rf.nextIndex[server] = max(args.LastIncludedIndex+1, rf.nextIndex[server])
	// Debug(dLeader, "S%d T%d installSnapshotOnceTo from S%d, set nextIndex %d", rf.me, rf.CurrentTerm, server, rf.nextIndex[server])
	return true
}

func (rf *Raft) installSnapshotTo(server int) {
	for n_tried := 0; !rf.killed() && n_tried < N_MAX_TRY; n_tried++ {
		if rf.installSnapshotOnceTo(server) {
			return
		}
		time.Sleep(time.Duration(TIGHT_LOOP_INTERVAL) * time.Millisecond)
	}
}

// the service says it has created a snapshot that has
// all info up to and including index. this means the
// service no longer needs the log through (and including)
// that index. Raft should now trim its log as much as possible.
func (rf *Raft) Snapshot(index int, snapshot []byte) {
	// Your code here (3D).

	rf.mu.Lock()
	if relativeIndex := rf.abs2Relative(index); relativeIndex >= 1 {
		rf.updateSnapshot(index, rf.Log[relativeIndex].Term, snapshot)
		// Debug(dLog, "S%d T%d, application snapshot index %d, term %d", rf.me, rf.CurrentTerm, index, rf.Log[0].Term)
	}
	rf.mu.Unlock()
}
